import numpy as np
import pandas as pd
from talib import abstract as taabstract
from jili.core import get_models
from jili.calc.ranged import analysis_rangeds
import talib
import copyreg
import pickle
ta_cmds=talib.get_functions()
class selfta:
    def __init__(self):
        self.calcors=get_models("jili.calc.ta",path_depth=2,model_type="class",included__init__=False)
        self.calcors_names=list(self.calcors.keys())
    def creact(self,ta_arg):
        cmd=ta_arg["cmd"]
        call= self.calcors.get(cmd,None)
        ta=None
        if call:
            ta=call(ta_arg)
        return ta
selftas=selfta()
def get_ma_batch(period,matype=0):
    if matype in [0,1,2,5,9]:
        return period
    elif matype ==3:
        return 2*period
    elif matype ==4:
        return 3*period
    elif matype ==8:
        return 5*period
    elif matype ==7:
        return 7*period
    elif matype ==6:
        return period+1
    else:
        return period
def creat_ta_calcors(ta_args):
    c=[]
    for i in ta_args:
        cmd=i["cmd"]
        if cmd in ta_cmds:
            c.append(ta(i))
        elif cmd in selftas.calcors_names:
            c0=selftas.creact(i)
            if c0:
                c.append(c0)
        elif cmd in analysis_rangeds.calcors_names:
            c0 = analysis_rangeds.creact(i)
            if c0:
                c.append(c0)
        else:
            c.append(ta(i))
    return c
def creat_ta_calcor(ta_arg,obj=""):
    cmd = ta_arg["cmd"]
    i=ta_arg
    if cmd in ta_cmds:
        isself=False
        if i["cmd"]=="MA" and "arg" in i.keys():
            if "matype" in i["arg"].keys():
                if i["arg"]["matype"] in [9,10]:
                    isself=True
        if isself:
            if i["arg"]["matype"]==9:
                i["cmd"]="LLT"
            else:
                i["cmd"]="ICU"
            return selftas.creact(i)
        else:
            return ta(i,obj)
    elif cmd in selftas.calcors_names:
        return selftas.creact(i)
    elif cmd in analysis_rangeds.calcors_names:
        return analysis_rangeds.creact(i)
    else:
        t0=ta()
        t0.setarg(i,obj)
        return t0
class ta:
    """:param
    1.只要求数据有input定义即可,不要求遵循k线原则，
    2.入参是dict 不支持list等
    3.支持复杂dict数据接口  可以深度搜索数据，结果也可以放在这个dict的层级字段中
    4.层级是用 A.B.C  data[A][B][C]
    5.同一个数据 以时间timekey为关键字更新 重复的则重复更新
    """
    def __init__(self):
        self.batch = 1
        self.lasttimekey = None
        self.input=[]
        # self.input = ["open","close","high","low","volume"]
        self.out = []
        self.hisbars = {}
        self.isoneout = True

    def reset(self):
        for i in self.input:
            self.hisbars[i]=[]
        self.init_ta(isforce=True)
    def drop_ta(self):
        self.ta=None
    def init_ta(self,isforce=False):
        if self.ta is None:
            self.ta = taabstract.Function(self.cmd)
        elif isforce:
            self.ta = taabstract.Function(self.cmd)
        if "input" in self.ta_arg.keys():
            input_names = self.ta_arg["input"]
            self.ta.set_input_names(input_names)
        if "arg" in self.ta_arg.keys():
            self.ta.set_parameters(self.ta_arg["arg"])
        self.batch=self.ta.lookback+1

    # 自定义序列化方法
    # def serialize(self):
    #     state = self.__dict__.copy()
    #     state["ta"]=None
    #     # Remove the unpicklable entries.
    #     # state['ta']=None
    #     print(state)
    #     return state

    # 自定义反序列化方法

    # def deserialize(self,state):
    #     # state = pickle.loads(serialized_str)
    #     if state["ta"]   is None:
    #         state["ta"] = taabstract.Function(state["cmd"])
    #     if "input" in state["ta_arg"].keys():
    #         input_names = state["ta_arg"]["input"]
    #         state["ta"].set_input_names(input_names)
    #     if "arg" in state["ta_arg"].keys():
    #         state["ta"].set_parameters(state["ta_arg"]["arg"])
    #     self.__dict__.update(state)
    #     return self
    # def __serialize__(self):
    #     """自定义序列化方法"""
    #     state = self.__getstate__()
    #
    #     return state
    #
    # def __deserialize__(self, data):
    #     """自定义反序列化方法"""
    #     self.__setstate__(data)
    def __reduce__(self):
        """实现该魔术方法以使用自定义序列化方法"""
        # print("pickle_ta:")
        state = self.__dict__.copy()
        if talib.__version__ > "0.4.24":
            del state["ta"]
        # print("__reduce__:", state)
        return ta, (state["ta_arg"], "",), state
    def __getstate__(self):
        # print("__getstate__")
        # Copy the object's state from self.__dict__ which contains
        # all our instance attributes. Always use the dict.copy()
        # method to avoid modifying the original state.
        state = self.__dict__.copy()
        if talib.__version__ > "0.4.24":
            del state["ta"]
        # Remove the unpicklable entries.
        # state['ta']=None
        # print(state)
        # print("__getstate__",state)
        return state

    def __setstate__(self, state):
        # Restore instance attributes (i.e., filename and lineno).
        # print("__setstate__",state)
        # if state["ta"]   is None:
        #     state["ta"] = taabstract.Function(state["cmd"])
        # if "input" in state["ta_arg"].keys():
        #     input_names = state["ta_arg"]["input"]
        #     state["ta"].set_input_names(input_names)
        # if "arg" in state["ta_arg"].keys():
        #     state["ta"].set_parameters(state["ta_arg"]["arg"])
        self.__dict__.update(state)
        # Restore the previously opened file's state. To do so, we need to
        # reopen it and read from it until the line count is restored.
        if talib.__version__ > "0.4.24":
            if self.ta   is None:
                cmd=self.cmd
                self.ta = taabstract.Function(cmd)
            if "input" in self.ta_arg.keys():
                input_names = self.ta_arg["input"]
                self.ta.set_input_names(input_names)
                # print(input_names)
            if "arg" in self.ta_arg.keys():
                self.ta.set_parameters(self.ta_arg["arg"])
                # print(self.ta_arg["arg"])
            self.batch = self.ta.lookback + 1
        # print(self.ta)
    def setarg(self, ta_arg,obj=""):
        self.ta_arg = ta_arg
        self.obj = obj
        self.outall = ta_arg.get("outall", True)
        self.default = ta_arg.get("default", None)
        self.isfirst = True
        self.ta = None
        if "cmd" in ta_arg.keys():
            cmd = ta_arg["cmd"]
            self.ta = taabstract.Function(cmd)
            self.cmd = cmd
            if "input" in ta_arg.keys():
                input_names = ta_arg["input"]
                self.ta.set_input_names(input_names)
            else:
                input_names = self.ta.get_input_names()
            # if "input" in ta.keys():
            #     input = ta["input"]
            #     fun.set_input_names(input)
            if isinstance(input_names, list):
                for i in input_names:
                    if i not in self.input:
                        self.input.append(i)
            else:
                for i in input_names.values():
                    if isinstance(i, list):
                        for j in i:
                            if i not in self.input:
                                self.input.append(j)
                    else:
                        if i not in self.input:
                            self.input.append(i)
            # if "batch" in ta_arg.keys():
            #     self.batch = ta_arg["batch"]
            if "fillna" in ta_arg.keys():
                self.fillna = ta_arg["fillna"]
            else:
                self.fillna = None
            if "arg" in ta_arg.keys():
                self.ta.set_parameters(ta_arg["arg"])
                if "timeperiod" in ta_arg["arg"].keys():
                    if ta_arg["arg"]["timeperiod"] > self.batch:
                        self.batch = ta_arg["arg"]["timeperiod"]
            self.batch = self.ta.lookback + 1
            if "out" in ta_arg.keys():
                self.out = ta_arg["out"]
                if len(self.out) > 1:
                    self.isoneout = False
            else:
                for i in self.call.output_names:
                    if i == "real":
                        self.out.append(cmd)
                    else:
                        self.out.append(cmd + "." + i)
                if len(self.out) > 1:
                    self.isoneout = False
        else:
            print("err  no cmd can not init talib")
    def sethisbar(self, bars):
        if isinstance(bars, list):
            self.lasttimekey = bars[-1]["timekey"]
            hisbars = pd.DataFrame(bars)
            for i in self.input:
                self.hisbars[i] = np.array(hisbars[i], dtype=np.double)
        elif isinstance(bars, pd.DataFrame):
            self.lasttimekey = bars["timekey"][-1].to_pydatetime()
            for i in self.input:
                self.hisbars[i] = np.array(bars[i], dtype=np.double)
    def getdata(self,data,key):
        if "." in key:
            for i in key.split("."):
                data=data[i]
                if data   is None:
                    data = np.nan
            return data
        else:
            if data[key]  is None:
                return np.nan
            else:
                return data[key]
    def setdata(self,data,t):
        for k, v in t.items():
            if "." in k:
                d = data
                l = k.split(".")
                n = len(l)
                nn = 0
                for i in l:
                    nn = nn + 1
                    if nn < n:
                        if i not in d.keys():
                            d[i] = {}
                        d = d[i]
                    else:
                        d[i] = v
            else:
                data[k] = v
        return data
    def ondata(self,data,timekey=None):#支持多层级数据检索与
        if self.lasttimekey:
            if timekey:
                tk=timekey
            else:
                tk = data["timekey"]
            if tk == self.lasttimekey:
                for i in self.input:
                    data0 = self.hisbars[i]
                    data0[-1] = self.getdata(data,i)
            else:
                self.lasttimekey = tk
                for i in self.input:
                    n = len(self.hisbars[i])
                    data00=self.getdata(data,i)
                    if self.batch > n:
                        self.hisbars[i].append(data00)
                    else:
                        self.hisbars[i].pop(0)
                        self.hisbars[i].append(data00)
        else:
            if timekey:
                self.lasttimekey=timekey
            else:
                self.lasttimekey = data["timekey"]
            for i in self.input:
                data00 = self.getdata(data, i)
                self.hisbars[i] = [data00]
        input = self.input[0]
        n = len(self.hisbars[input])
        if n >= self.batch:
            isrun = True
        else:
            isrun = False
        if isrun and self.isfirst:
            for i in self.hisbars.values():
                if pd.isna(i[-1]):
                    isrun = False
                    break
            if isrun:
                self.isfirst = False
        if isrun:
            hisbars={}
            for i,v in self.hisbars.items():
                hisbars[i]=np.array(v, dtype=np.double)
            self.ta.set_input_arrays(hisbars)
            if self.isoneout:
                try:
                    t = {self.out[0]: self.ta.outputs[-1]}
                except Exception as e:
                    print(e)
            else:
                n = 0
                t = {}
                for i in self.out:
                    if i != "_":
                        t[i] = self.ta.outputs[n][-1]
                    n = n + 1
        else:
            t = {}
            for i in self.out:
                if i != "_":
                    t[i] = self.default
        if self.outall:
            return self.setdata(data, t)
        else:
            return t
    def get_fillna(self,a,data):
        if pd.isna(a):
            if self.fillna is None:
                if data:
                    return data[-1]
            else:
                return self.fillna
        return a
    def onbar(self, bar):
        if self.lasttimekey:
            tk = bar["timekey"]
            if tk == self.lasttimekey:
                for i in self.input:
                    data = self.hisbars[i]
                    d0=self.get_fillna(bar[i],data)
                    data[-1] = d0
            else:
                self.lasttimekey = tk
                for i in self.input:
                    n = len(self.hisbars[i])
                    d0 = self.get_fillna(bar[i], self.hisbars[i])
                    if self.batch > n:
                        self.hisbars[i].append(d0)
                    else:
                        self.hisbars[i].pop(0)
                        self.hisbars[i].append(d0)
        else:
            self.lasttimekey = bar["timekey"]
            for i in self.input:
                d0 = bar[i]
                self.hisbars[i] = [d0]
        input=self.input[0]
        n = len(self.hisbars[input])
        if n >= self.batch:
            isrun=True
        else:
            isrun = False
        if isrun and self.isfirst:
            for i in self.hisbars.values():
                if pd.isna(i[-1]):
                    isrun=False
                    break
            if isrun:
                self.isfirst=False
        if isrun:
            hisbars = {}
            for i, v in self.hisbars.items():
                hisbars[i] = np.array(v, dtype=np.double)
            self.ta.set_input_arrays(hisbars)
            if self.isoneout:
                # t={self.out[0]: self.ta.outputs[-1]}
                try:
                    t={self.out[0]: self.ta.outputs[-1]}
                except Exception as e:
                    print(e)
                    t = {self.out[0]: self.default}
            else:
                n = 0
                t = {}
                for i in self.out:
                    if i !="_":
                        t[i] = self.ta.outputs[n][-1]
                    n = n + 1
        else:
            t = {}
            for i in self.out:
                if i != "_":
                    t[i] = self.default
        if self.outall:
            bar.update(t)
            return bar
        else:
            return t
def pickle_ta(c):
    # print("pickle_ta:")
    state = c.__dict__.copy()
    del state["ta"]
    print("pickle_ta:",state)
    return ta, (state["ta_arg"],"",),state
# copyreg.pickle(ta, pickle_ta)
if __name__ == "__main__":
    import datetime
    a={'a': {'b': 0.1}, 'c': 1, 'timekey': datetime.datetime(2022, 8, 22, 0, 0)}
    cos = ta({'calc_cmd': 'ta', 'cmd': 'cos', 'out': ['cos.cos'], 'input': {'price': 'a.b'}, 'arg': {}, 'batch': 1,"outall":True})
    print(cos.ondata(a))